package com.github.grpc.server;

import io.grpc.BindableService;
import io.grpc.Server;
import io.grpc.ServerBuilder;
import io.grpc.ServerServiceDefinition;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.springframework.aop.framework.AopProxyUtils;
import org.springframework.aop.scope.ScopedProxyUtils;
import org.springframework.beans.BeansException;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.context.SmartLifecycle;
import org.springframework.util.CollectionUtils;

import java.io.IOException;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.stream.Collectors;

@Slf4j
public class GrpcServer implements ApplicationContextAware, SmartLifecycle {
    private final int port;
    private final ServerBuilder serverBuilder;

    private Server server;
    private final AtomicBoolean running = new AtomicBoolean(false);
    private ApplicationContext applicationContext;

    public GrpcServer(int port) {
        this(port, ServerBuilder.forPort(port));
    }

    public GrpcServer(int port, ServerBuilder serverBuilder) {
        this.port = port;
        this.serverBuilder = serverBuilder;
    }

    @Override
    public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
        this.applicationContext = applicationContext;
    }

    @Override
    @SneakyThrows(IOException.class)
    public void start() {
        this.server = serverBuilder.addServices(getGrpcServices()).build();
        this.server.start();
        this.running.set(true);
        log.info("Grpc server started, listening on {}.", port);
    }

    private List<ServerServiceDefinition> getGrpcServices() {
        Map<String, Object> beans = this.applicationContext.getBeansWithAnnotation(GrpcService.class)
                // ScopedProxyUtils.isScopedTarget：如果Scope的ScopedProxyMode设置未非Default和ON
                // 原始对象的名字会被加上"scopedTarget."的前缀，而通过类型或者beanName查找到的都是代理对象
                // 此处过滤掉原始对象，使用代理对象
                .entrySet().stream().filter(entry -> !ScopedProxyUtils.isScopedTarget(entry.getKey()))
                .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
        if (CollectionUtils.isEmpty(beans)) {
            return Collections.EMPTY_LIST;
        }

        return beans.entrySet().stream().map(entry -> {
            Object bean = entry.getValue();
            Class<?> clazz = AopProxyUtils.ultimateTargetClass(bean);
            if (!BindableService.class.isAssignableFrom(bean.getClass())) {
                throw new IllegalStateException(clazz + "is not instance of " + BindableService.class.getName());
            }

            return ((BindableService) bean).bindService();
        }).collect(Collectors.toList());
    }

    @Override
    @SneakyThrows(InterruptedException.class)
    public void stop() {
        running.set(false);
        log.info("*** shutting down gRPC server since JVM is shutting down.");
        if (server != null) {
            server.shutdown().awaitTermination(30, TimeUnit.SECONDS);
        }
        log.info("*** server shut down");
    }

    @Override
    public boolean isRunning() {
        return this.running.get();
    }
}
